import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import matplotlib.pyplot as plt
import os
import joblib


def plot_history(val_scores, fold_count):
    """
    绘制每一折的 MSE 指标（英文标注）
    """
    plt.figure(figsize=(10, 5))
    plt.plot(val_scores, label='Validation MSE', marker='o')
    plt.title('K-Fold Validation MSE Scores')
    plt.xlabel('Fold')
    plt.ylabel('MSE')
    plt.legend()
    plt.grid(True)

    # 创建输出目录并保存图像
    output_dir = "../data/output/"
    os.makedirs(output_dir, exist_ok=True)
    plt.savefig(os.path.join(output_dir, "rf_training_history.png"))
    plt.close()  # 非交互环境避免弹窗显示图像


def train_model(X, y, n_splits=5, random_state=42):
    """
    使用 KFold 交叉验证训练模型，并返回每折的评估结果和模型
    """
    kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    fold_no = 1
    val_scores = []
    mae_scores = []
    r2_scores = []

    for train_index, val_index in kf.split(X):
        print(f"\n 开始第 {fold_no} 折交叉验证")
        X_train, X_val = X[train_index], X[val_index]
        y_train, y_val = y[train_index], y[val_index]

        # 构建模型
        model = RandomForestRegressor(
            n_estimators=100,
            max_depth=None,
            min_samples_split=2,
            min_samples_leaf=1,
            n_jobs=-1,
            random_state=random_state
        )

        # 训练模型
        model.fit(X_train, y_train)

        # 预测与评估
        y_pred = model.predict(X_val)
        mse = mean_squared_error(y_val, y_pred)
        mae = mean_absolute_error(y_val, y_pred)
        r2 = r2_score(y_val, y_pred)

        val_scores.append(mse)
        mae_scores.append(mae)
        r2_scores.append(r2)

        print(f" 第 {fold_no} 折验证完成 - Val MSE: {mse:.6f}, MAE: {mae:.6f}, R²: {r2:.6f}")

        # 保存每折的预测结果（可选）
        result_df = pd.DataFrame({
            "Fold": [fold_no] * len(y_val),
            "True_Value": y_val,
            "Predicted_Value": y_pred
        })
        result_path = "../data/output/kfold_predictions.parquet"
        if fold_no == 1:
            result_df.to_parquet(result_path, engine='pyarrow', index=False)
        else:
            prev_df = pd.read_parquet(result_path)
            combined_df = pd.concat([prev_df, result_df], ignore_index=True)
            combined_df.to_parquet(result_path, engine='pyarrow', index=False)

        fold_no += 1

    # 输出平均指标
    avg_mse = np.mean(val_scores)
    avg_mae = np.mean(mae_scores)
    avg_r2 = np.mean(r2_scores)

    print("\n 五折平均评估指标：")
    print(f" 平均 MSE: {avg_mse:.6f}")
    print(f" 平均 MAE: {avg_mae:.6f}")
    print(f" 平均 R² Score: {avg_r2:.6f}")

    return model, avg_mse, avg_mae, avg_r2


if __name__ == "__main__":
    print(" 正在加载数据...")
    df = pd.read_parquet("../data/features/feature_engineered_data.parquet")
    X = df.drop(columns=['Average_Fare']).values
    y = df['Average_Fare'].values

    print(f" 数据形状：X={X.shape}, y={y.shape}")

    # 训练模型
    best_model, avg_mse, avg_mae, avg_r2 = train_model(X, y)

    # 保存模型
    os.makedirs("../data/models", exist_ok=True)
    model_path = "../data/models/random_forest_kfold.pkl"
    joblib.dump(best_model, model_path)
    print(f" 已保存 KFold 最佳模型到 {model_path}")

    # 可视化 MSE 曲线
    plot_history(avg_mse, fold_count=5)
